import numpy as np
import cv2
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import logging


ax = plt.gca()

logger = logging.getLogger("ssjhs")

# 168,127
# 168,87

# import numpy as np
# from scipy.optimize import curve_fit

# def func(X, a, b, c):
#     x,y = X
#     return np.log(a) + b*np.log(x) + c*np.log(y)

# # some artificially noisy data to fit
# x = np.linspace(0.1,1.1,101)
# y = np.linspace(1.,2., 101)
# a, b, c = 10., 4., 6.
# z = func((x,y), a, b, c) * 1 + np.random.random(101) / 100

# # initial guesses for a,b,c:
# p0 = 8., 2., 7.
# coeff = curve_fit(func, (x,y), z, p0)


def __getRoi(image, cx, cy, w, h):
    height, width = image.shape[0:2]
    x0 = cx - w if cx - w > 0 else 0
    x1 = cx + w if cx + w < width else width - 1

    y0 = cy - h if cy - h > 0 else 0
    y1 = cy + h if cy + h < height else height - 1

    return x0, x1, y0, y1


def correct(X, k1, k2, k3, p1, p2):
    x, y = X
    r2 = x**2 + y**2

    # 径向畸变 radial distortion
    x_rd = x * (1 + k1 * r2 + k2 * (r2**2) + k3 * (r2**3))
    y_rd = y * (1 + k1 * r2 + k2 * (r2**2) + k3 * (r2**3))

    # 切向畸变 tangential distortion
    x_td = 2 * p1 * x * y + p2 * (r2 + 2 * (x**2))
    y_td = 2 * p2 * x * y + p1 * (r2 + 2 * (y**2))

    return x_rd + x_td, y_rd + y_td


def correctXY(X, k1, k2, k3, p1, p2):
    x, y = X
    r2 = x**2 + y**2

    # 径向畸变 radial distortion
    x_rd = x * (1 + k1 * r2 + k2 * (r2**2) + k3 * (r2**3))
    y_rd = y * (1 + k1 * r2 + k2 * (r2**2) + k3 * (r2**3))

    # 切向畸变 tangential distortion
    x_td = 2 * p1 * x * y + p2 * (r2 + 2 * (x**2))
    y_td = 2 * p2 * x * y + p1 * (r2 + 2 * (y**2))

    return x_rd + x_td


def correctX(X, k1, k2, k3, p1, p2):
    return correct(X, k1, k2, k3, p1, p2)[0]


def correctY(X, k1, k2, k3, p1, p2):
    return correct(X, k1, k2, k3, p1, p2)[1]


def findCircleCenter(img, CX, CY, D=10, R=8):
    global origin

    x, y = CX, CY
    # cv2.rectangle(origin, (x-D, y-D), (x+D, y+D), (0, 0, 255))
    # cv2.imshow('origin', origin)
    # cv2.waitKey()

    array = np.empty(0)
    for i in range(-D, D):
        x = CX + i
        x0, x1, y0, y1 = __getRoi(img, x, y, D, D)
        roi = img[y0:y1, x0:x1]
        # roi = img[y - D : y + D, x - D : x + D]
        aver = np.average(roi, axis=0)
        array = np.append(array, np.average(aver[D - R : D + R]))

        # cv2.rectangle(origin, (x-D, y-D), (x+D, y+D), (0, 0, 255))
        # cv2.imshow('origin', origin)
        # cv2.waitKey()
        # cv2.destroyAllWindows()

        # plt.plot(aver)
        # plt.show()
    x = np.argmax(array)
    x = CX - D + x

    array = np.empty(0)
    for i in range(-D, D):
        y = CY + i
        x0, x1, y0, y1 = __getRoi(img, x, y, D, D)
        roi = img[y0:y1, x0:x1]
        # roi = img[y - D : y + D, x - D : x + D]
        aver = np.average(roi, axis=1)
        array = np.append(array, np.average(aver[D - R : D + R]))

        # cv2.rectangle(origin, (x-D, y-D), (x+D, y+D), (0, 0, 255))
        # cv2.imshow('origin', origin)
        # cv2.waitKey()
        # cv2.destroyAllWindows()

        # plt.plot(aver)
        # plt.show()
    y = np.argmax(array)
    y = CY - D + y

    roi = img[y - D : y + D, x - D : x + D]

    # calculate moments of binary image
    M = cv2.moments(roi)
    if M["m00"] == 0:
        return None

    # calculate x,y coordinate of center
    cx = (M["m10"] / M["m00"]) + x - D
    cy = (M["m01"] / M["m00"]) + y - D

    return cx, cy


def calibrate(img, D=40, N=5):
    MID = int((N - 1) / 2)

    global origin
    origin = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)

    CY, CX = img.shape
    CY, CX = round(CY / 2), round(CX / 2)
    roi = img.copy()
    # cv2.imwrite('origin.jpg', roi)

    ###############################################################
    # 修正亮度
    r = 40
    roi[:r] = np.where((roi[:r] > 30), 255, roi[:r])
    # cv2.imshow('roi', roi)
    # cv2.waitKey()

    y, x = np.ogrid[0 : img.shape[0], 0 : img.shape[1]]

    r = 90
    mask = (x - CX) ** 2 + (y - CY) ** 2 < r**2
    roi[mask] = np.where((roi[mask] < 70), 0, roi[mask])

    r = 100
    mask = (x - CX) ** 2 + (y - CY) ** 2 > r**2
    roi[mask] = np.where((roi[mask] > 30), 255, roi[mask])

    r = 100
    mask = (x - CX) ** 2 + (y - CY) ** 2 < r**2
    roi[mask] = np.where((roi[mask] < 50), 0, roi[mask])

    # cv2.imshow("roi", roi)
    # cv2.waitKey()

    ret, thresh = cv2.threshold(roi, 60, 255, cv2.THRESH_BINARY_INV)
    # cv2.imshow("thresh", thresh)
    # cv2.waitKey()

    ###############################################################
    # 找中心点
    cx_float, cy_float = findCircleCenter(thresh, CX, CY, D=20)
    cx, cy = round(cx_float), round(cy_float)
    cv2.circle(origin, (cx, cy), 1, (20, 20, 220), 2)
    logger.info(f"cp: {cx}, {cy}.  b:{img[cy][cx]} b:{img[217][5]}")

    ###############################################################
    # 生成矩阵 5 x 5
    matrix = []
    for y in range(cy - D * MID, cy + D * (MID + 1), D):
        row = []
        for x in range(cx - D * MID, cx + D * (MID + 1), D):
            result = findCircleCenter(thresh, x, y)
            if result is None:
                raise Exception("findCircleCenter is None")
            row.append(result)
            # cv2.circle(origin, (x, y), 1, (20, 220, 20), 2)
            # cv2.imshow("origin", origin)
            # cv2.waitKey()
        matrix.append(row)

    cv2.circle(origin, (CX, CY), 1, (120, 100, 220), 2)

    ###############################################################
    # 计算点间距
    x1, y1 = matrix[MID][MID - 1]
    x2, y2 = matrix[MID][MID + 1]
    rx = (np.sqrt(abs(x2 - x1) ** 2 + abs(y2 - y1) ** 2)) / 2

    x1, y1 = matrix[MID - 1][MID]
    x2, y2 = matrix[MID + 1][MID]
    ry = (np.sqrt(abs(x2 - x1) ** 2 + abs(y2 - y1) ** 2)) / 2

    dx = (matrix[MID + 1][MID][0] - matrix[MID - 1][MID][0]) / 2
    dy = (matrix[MID][MID + 1][1] - matrix[MID][MID - 1][1]) / 2

    # x,y = matrix[MID][MID][0], matrix[MID][MID][1]
    # cv2.circle(origin, (round(x),round(y)), 4, (220, 222, 20), 2)
    # x,y = matrix[MID][MID+1][0], matrix[MID][MID+1][1]
    # cv2.circle(origin, (round(x),round(y)), 4, (220, 222, 20), 2)

    print(rx, ry, dx, dy)

    ###############################################################
    # 拟合修正
    x_origin, y_origin, x_correct, y_correct = [], [], [], []
    input, output = [], []
    for i in range(N):
        for j in range(N):
            x, y = matrix[i][j]
            input.append((x, y))
            input.append((y, x))
            output.append(cx + (j - MID) * rx)
            output.append(cy + (i - MID) * ry)

            x_origin.append(x)
            y_origin.append(y)

            x_cor = round(cx + (j - MID) * D)
            y_cor = round(cy + (i - MID) * D)
            x_correct.append(x_cor)
            y_correct.append(y_cor)
            # cv2.circle(origin, (round(x), round(y)), 1, (20, 220, 20), 2)
            cv2.circle(origin, (x_cor, y_cor), 1, (220, 20, 20), 2)

    p0 = [0, 0, 0, 0, 0]
    # print(correct((1,2),1,2,3,4,5))
    # print(correct((2,1),1,2,3,4,5))
    input = np.array(input)

    popt, pcov = curve_fit(correctXY, (input[:, 0], input[:, 1]), output, p0)
    coeff = popt

    p0 = [0, 0, 0, 0, 0]
    popt, pcov = curve_fit(correctX, (x_origin, y_origin), x_correct, p0)
    coeff_x = popt

    p0 = [0, 0, 0, 0, 0]
    popt, pcov = curve_fit(correctY, (x_origin, y_origin), y_correct, p0)
    coeff_y = popt

    # np.set_printoptions(precision=3)
    print(coeff_x)
    print(coeff_y)
    print(correctX([234,123], *coeff_x))
    print(correctY([234,123], *coeff_y))

    for i in range(N):
        for j in range(N):
            # x, y = correct(matrix[i][j], *coeff)
            x = correctX(matrix[i][j], *coeff_x)
            y = correctY(matrix[i][j], *coeff_y)

            ax.scatter(x, y, c="g")

            x, y = round(x), round(y)
            cv2.circle(origin, (x, y), 1, (20, 20, 220), 2)

            # print(x, y)

            # x = output[(i * N + j) * 2]
            # y = output[(i * N + j) * 2 + 1]
            # x, y = round(x), round(y)
            # cv2.circle(origin, (x, y), 1, (220, 20, 20), 2)

    ###############################################################
    # ax.invert_yaxis()
    # plt.show()

    # plt.axline(matrix[2][2], matrix[2][3])

    # # some artificially noisy data to fit
    # x = np.linspace(0.1,1.1,101)
    # y = np.linspace(1.,2., 101)
    # a, b, c = 10., 4., 6.
    # z = func((x,y), a, b, c) * 1 + np.random.random(101) / 100

    # # initial guesses for a,b,c:
    # p0 = 8., 2., 7.
    # print curve_fit(func, (x,y), z, p0)

    # for x, y in matrix[1]:
    #     cv2.circle(origin, (x, y), 1, (20, 220, 20), 2)

    cv2.imshow("origin", origin)
    # plt.show()
    cv2.waitKey()

    # cy = CY
    # for i in range(-d, d + 1):
    #     cx = CX + i
    #     roi = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    #     cv2.rectangle(roi, (cx - d, cy - d), (cx + d, cy + d), (0, 0, 255))

    #     cv2.imshow("roi", roi)
    #     cv2.waitKey()


# img = cv2.imread('1.png')
# gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# roi = gray

# # convert the grayscale image to binary image
# ret, thresh = cv2.threshold(roi, 127, 255, 0)

# contours, hierarchy = cv2.findContours(
#     thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)


# for i, cnt in enumerate(contours):
#     cv2.drawContours(img, contours, i, (0, 255, 0), 1)

# # calculate moments of binary image
# M = cv2.moments(thresh)

# # calculate x,y coordinate of center
# cX = int(M["m10"] / M["m00"])
# cY = int(M["m01"] / M["m00"])

# # put text and highlight the center
# cv2.circle(img, (cX, cY), 5, (255, 255, 255), -1)
# cv2.putText(img, "centroid", (cX - 25, cY - 25),
#             cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)

# cv2.imshow('roi', roi)
# cv2.imshow('img', img)
# cv2.waitKey()

# cy = CY
# for i in range(-d, d+1):
#     cx = CX + i
#     roi = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
#     cv2.rectangle(roi, (cx-d, cy-d), (cx+d, cy+d), (0, 0, 255))

#     cv2.imshow('roi', roi)
#     cv2.waitKey()
